import numpy as np
import torch
import matplotlib.pyplot as plt


def fill_cylinder_with_particles(inner_radius, outer_radius, length_x, spacing, random_disturbance_mag=0.8):
    """Fill hollow cylinder with particles"""
    
    n_particles_x = int(length_x / spacing)
    n_particles_theta = int(2 * np.pi * outer_radius / spacing)
    n_particles_r = int((outer_radius - inner_radius) / spacing)
    
    particles = []
    
    x_positions = np.linspace(spacing/2, length_x - spacing/2, n_particles_x)
    r_positions = np.linspace(inner_radius + spacing/2, outer_radius - spacing/2, n_particles_r)
    theta_positions = np.linspace(0, 2*np.pi - 2*np.pi/n_particles_theta, n_particles_theta)
    
    for x in x_positions:
        for r in r_positions:
            for theta in theta_positions:
                y = r * np.cos(theta)
                z = r * np.sin(theta)
                particles.append([x, y, z])
    
    particles = np.array(particles)
    
    # Add noise
    disturbance = np.random.uniform(
        -spacing/2 * random_disturbance_mag,
        spacing/2 * random_disturbance_mag,
        size=particles.shape
    )
    
    disturbed_particles = particles + disturbance
    torch_particles = torch.tensor(disturbed_particles, dtype=torch.float32, requires_grad=True)
    
    return torch_particles


def fill_cylinder_end_caps(inner_radius, outer_radius, spacing, x_position, random_disturbance_mag=0.8):
    """Generate end cap particles"""
    
    particles = []
    n_particles_r = int((outer_radius - inner_radius) / spacing)
    
    for i in range(n_particles_r):
        r = inner_radius + (i + 0.5) * spacing
        circumference = 2 * np.pi * r
        n_particles_theta = max(1, int(circumference / spacing))
        
        for j in range(n_particles_theta):
            theta = 2 * np.pi * j / n_particles_theta
            y = r * np.cos(theta)
            z = r * np.sin(theta)
            particles.append([x_position, y, z])
    
    particles = np.array(particles)
    
    disturbance = np.random.uniform(
        -spacing/2 * random_disturbance_mag,
        spacing/2 * random_disturbance_mag,
        size=particles.shape
    )
    
    disturbed_particles = particles + disturbance
    torch_particles = torch.tensor(disturbed_particles, dtype=torch.float32, requires_grad=True)
    
    return torch_particles


def create_cylinder_drum(inner_radius, outer_radius, length_x, spacing, include_end_caps=True, random_disturbance_mag=0.8):
    """Create complete cylinder drum geometry"""
    
    wall_particles = fill_cylinder_with_particles(
        inner_radius, outer_radius, length_x, spacing, random_disturbance_mag
    )
    
    all_particles = [wall_particles]
    
    if include_end_caps:
        left_cap = fill_cylinder_end_caps(
            inner_radius, outer_radius, spacing, 0, random_disturbance_mag
        )
        all_particles.append(left_cap)
        
        right_cap = fill_cylinder_end_caps(
            inner_radius, outer_radius, spacing, length_x, random_disturbance_mag
        )
        all_particles.append(right_cap)
    
    cylinder_particles = torch.vstack(all_particles)
    return cylinder_particles


def visualize_cylinder_particles(particles_array, title="Cylinder Particles"):
    """Basic 3D visualization"""
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    ax.scatter(particles_array[:, 0], particles_array[:, 1], particles_array[:, 2], 
               s=1, alpha=0.6)
    
    ax.set_xlabel('X')
    ax.set_ylabel('Y') 
    ax.set_zlabel('Z')
    ax.set_title(title)
    
    max_range = np.array([
        particles_array[:, 0].max() - particles_array[:, 0].min(),
        particles_array[:, 1].max() - particles_array[:, 1].min(),
        particles_array[:, 2].max() - particles_array[:, 2].min()
    ]).max() / 2.0
    
    mid_x = (particles_array[:, 0].max() + particles_array[:, 0].min()) * 0.5
    mid_y = (particles_array[:, 1].max() + particles_array[:, 1].min()) * 0.5  
    mid_z = (particles_array[:, 2].max() + particles_array[:, 2].min()) * 0.5
    
    ax.set_xlim(mid_x - max_range, mid_x + max_range)
    ax.set_ylim(mid_y - max_range, mid_y + max_range)
    ax.set_zlim(mid_z - max_range, mid_z + max_range)
    
    plt.tight_layout()
    plt.show()


def locate_cylinder_particles(cylinder_particles: torch.tensor, cylinder_center: torch.tensor, rotation_angle: torch.tensor = None):
    """Position cylinder at specified location with optional rotation"""
    moved_particles = cylinder_particles.clone()
    
    if rotation_angle is not None:
        cos_angle = torch.cos(rotation_angle)
        sin_angle = torch.sin(rotation_angle)
        
        y_new = moved_particles[:, 1] * cos_angle - moved_particles[:, 2] * sin_angle
        z_new = moved_particles[:, 1] * sin_angle + moved_particles[:, 2] * cos_angle
        
        moved_particles[:, 1] = y_new
        moved_particles[:, 2] = z_new
    
    moved_particles[:, 0] += cylinder_center[0]
    moved_particles[:, 1] += cylinder_center[1] 
    moved_particles[:, 2] += cylinder_center[2]
    
    return moved_particles 